import torch
import io
from PIL import Image
from app.utils.image_preprocess import preprocess_image
from app.models.moisture_model import MoistureModel
from app.models.pest_model import PestModel
from config.settings import CHECKPOINT_DIR
import logging
import os

logger = logging.getLogger(__name__)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 全局模型变量
moisture_model = None
pest_model = None

MOISTURE_CLASSES = ["extreme_drought", "mild_drought", "moderate", "waterlogged", "normal"]
PEST_CLASSES = ["powdery_mildew", "black_spot", "aphids", "red_spider_mite", "scale_insects", "none"]

def load_moisture_model():
    """懒加载水分模型"""
    global moisture_model
    if moisture_model is None:
        try:
            moisture_model = MoistureModel()
            checkpoint_path = CHECKPOINT_DIR / "moisture_best.pth"
            
            # 检查是否有训练好的模型
            if checkpoint_path.exists():
                moisture_model.load_state_dict(
                    torch.load(checkpoint_path, map_location=device)
                )
                logger.info("加载训练好的水分模型")
            else:
                # 使用预训练模型
                logger.info("使用预训练ResNet18模型进行水分诊断")
                # 保持预训练权重，只修改最后的分类层
                moisture_model.eval()
            
            moisture_model.to(device)
            logger.info("水分模型加载成功")
        except Exception as e:
            logger.error(f"水分模型加载失败: {str(e)}")
            raise
    return moisture_model

def load_pest_model():
    """懒加载病虫害模型"""
    global pest_model
    if pest_model is None:
        try:
            pest_model = PestModel()
            checkpoint_path = CHECKPOINT_DIR / "pest_best.pth"
            
            # 检查是否有训练好的模型
            if checkpoint_path.exists():
                pest_model.load_state_dict(
                    torch.load(checkpoint_path, map_location=device)
                )
                logger.info("加载训练好的病虫害模型")
            else:
                # 使用预训练模型
                logger.info("使用预训练ResNet18模型进行病虫害诊断")
                # 保持预训练权重，只修改最后的分类层
                pest_model.eval()
            
            pest_model.to(device)
            logger.info("病虫害模型加载成功")
        except Exception as e:
            logger.error(f"病虫害模型加载失败: {str(e)}")
            raise
    return pest_model

def run_inference(image_bytes: bytes, model_type: str) -> tuple[str, float]:
    """运行推理"""
    try:
        # 验证输入
        if not image_bytes:
            raise ValueError("图像数据为空")
        
        # 加载图像
        image = Image.open(io.BytesIO(image_bytes))
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # 预处理图像
        tensor = preprocess_image(image).unsqueeze(0).to(device)
        
        # 根据模型类型进行推理
        if model_type == "moisture":
            model = load_moisture_model()
            classes = MOISTURE_CLASSES
        elif model_type == "pest":
            model = load_pest_model()
            classes = PEST_CLASSES
        else:
            raise ValueError(f"不支持的模型类型: {model_type}")
        
        # 推理
        with torch.no_grad():
            outputs = model(tensor)
        
        # 计算概率和预测结果
        probs = torch.softmax(outputs, dim=1)
        conf, pred = torch.max(probs, 1)
        conf = conf.item()
        pred = pred.item()
        
        result = classes[pred]
        logger.info(f"{model_type} 推理结果: {result}, 置信度: {conf:.4f}")
        
        return result, conf
        
    except Exception as e:
        logger.error(f"推理过程中发生错误: {str(e)}")
        raise